import multiprocessing
import threading
import time
import redis
from flask import Flask, request
import json
from xhs_dm_spider_cs import *

process_info_modify_lock = threading.Lock()
SPIDER_RESULT_KEY = "key:spider:{user_id}"
SPIDER_LOCK = "lock:spider:{user_id}"
"""测试REDIS"""
REDIS = redis.StrictRedis(host='r-bp162522qrwjh5e7jhpd.redis.rds.aliyuncs.com', port=6379,password='123sdmmh568%$#vs*!',db=22)
# """正式redis"""
# REDIS = redis.StrictRedis(host='r-bp1d0lcwddmu57gok3.redis.rds.aliyuncs.com', port=6379,password='123sdHs45GH68%$#vs*!', db=15)

SPIDER_STOP_KEY = "stop_auth_status:{user_id}"
SPIDER_STOP_VALUE = b"stop"  # 外部取消
SPIDER_STOP_VALUE_PROCESS_DONE = b"process_done"  # 程序自动完成


def spider_function(user_id: str):
    res = XHS_XPIDER(str(user_id))
    # try:
    result_data = res.xhs_main_app()
    print(result_data)

    # result_data_json = {"code":200,"msg":"exe_success","method":"POST","result":result_data}
    return_value = result_data
    # except Exception:
    #     return_value = {"code": 501, "msg": "网络连接超时，请稍后再试!"}

    REDIS.set(SPIDER_RESULT_KEY.format(user_id=user_id), json.dumps(return_value))
    # REDIS.expire(SPIDER_RESULT_KEY.format(user_id=user_id), 300)
    REDIS.set(
        SPIDER_STOP_KEY.format(user_id=user_id), SPIDER_STOP_VALUE_PROCESS_DONE
    )
    SPIDER_RESULT_KEY.format(user_id=user_id)
    # REDIS.expire(SPIDER_RESULT_KEY.format(user_id=user_id), 300)
    return return_value


def read_from_redis(user_id: str):
    key = SPIDER_STOP_KEY.format(user_id=user_id)
    while True:
        value = REDIS.get(key)
        if value in [SPIDER_STOP_VALUE, SPIDER_STOP_VALUE_PROCESS_DONE]:
            REDIS.delete(key)
            if value == SPIDER_STOP_VALUE_PROCESS_DONE:
                print("程序已完成")
            else:
                print(f"检测到值为stop, 退出{key}")

                REDIS.set(
                    SPIDER_RESULT_KEY.format(user_id=user_id),
                    json.dumps({"code": 400, "msg": "已取消"}),
                )
            return
        time.sleep(1)


def process_function(user_id: str):
    # 创建爬取数据线程并设置为守护线程
    spider_thread = threading.Thread(target=spider_function, args=(user_id,))
    spider_thread.daemon = True

    # 创建读取redis线程
    read_thread = threading.Thread(target=read_from_redis, args=(user_id,))

    # 启动线程
    spider_thread.start()
    read_thread.start()

    # 等待读取redis线程完成
    read_thread.join()

    print("Main thread: Read thread has exited, shutting down.")
    REDIS.delete(SPIDER_LOCK.format(user_id=user_id))


app = Flask(__name__)


@app.route("/process", methods=["GET", "POST"])
def process_data():
    if request.method == "POST":
        # 爬取请求
        print("当前为POST请求")
        data = request.get_json()

        user_id = data["user_id"]
        print("当前访问用户ID：", user_id)

        if REDIS.get(SPIDER_STOP_KEY.format(user_id=user_id)) == SPIDER_STOP_VALUE:

            return {"code": 500, "msg": "值为stop, 请稍后再试!"}

        user_spider_result_key = SPIDER_RESULT_KEY.format(user_id=user_id)
        spider_lock = SPIDER_LOCK.format(user_id=user_id)
        with process_info_modify_lock:
            lock = REDIS.get(spider_lock)
            if lock is not None:
                REDIS.set('xhs_chrome_open:{}'.format(user_id),'1')
                REDIS.expire('xhs_chrome_open:{}'.format(user_id), 300)
                return {"code": 500, "msg": "该用户正在爬取数据, 请稍后再试!"}

        process = multiprocessing.Process(
            target=process_function, args=(user_id,)
        )
        with process_info_modify_lock:
            REDIS.set(spider_lock, "running")
        process.start()  # 启动进程
        process.join()  # 等待进程结束, 等待爬取结束

        result = REDIS.get(user_spider_result_key)
        if result is None:
            return {"code": 501, "msg": "未获取到数据"}
        else:
            return json.loads(result)


@app.route("/stop_process", methods=["POST"])
def stop_process():
    data = request.get_json()
    user_id = data["user_id"]
    print("当前访问用户ID：", user_id)

    user_spider_lock = SPIDER_LOCK.format(user_id=user_id)
    with process_info_modify_lock:
        lock = REDIS.get(user_spider_lock)
        if lock is not None:
            user_spider_stop_key = SPIDER_STOP_KEY.format(user_id=user_id)
            REDIS.set(user_spider_stop_key, SPIDER_STOP_VALUE)
            return {"code": 200, "msg": "已设置停止信号"}
        else:
            return {"code": 400, "msg": "该用户未在运行!"}


def clear_spider_locks():
    _locks = REDIS.keys(SPIDER_LOCK.format(user_id="*"))
    if _locks:
        REDIS.delete(*_locks)


if __name__ == "__main__":
    # app.run(debug=True)
    clear_spider_locks()
    app.run(host="0.0.0.0", port=5009, threaded=True)
